import siot
import time
import socket
import cv2
import threading
import numpy as np


def preproc(raw_frame, thr=128):
    h, w, _ = raw_frame.shape
    raw = cv2.cvtColor(raw_frame[h // 2 - 80: h // 2 + 80,
                    w // 2 - 120: w // 2 + 120, :], cv2.COLOR_BGR2GRAY)
    cv2.imshow('Capture', raw)
    n = 1
    thr = np.mean(raw) * 1
    frame = cv2.filter2D(raw, cv2.CV_32F, np.array([[0, -n, 0], [-n, 4 * n + 1, -n], [0, -n, 0]], dtype=np.float32))
    frame = cv2.convertScaleAbs(frame)
    _, frame = cv2.threshold(frame, thr, 255, cv2.THRESH_BINARY_INV)
    frame = cv2.medianBlur(frame, 11)
    assert frame.shape == (160, 240)
    return frame, min(70, raw[80, 2])


isRunning = True


class CameraBufferCleanerThread(threading.Thread):
    def __init__(self, camera, name='camera-buffer-cleaner-thread'):
        self.camera = camera
        self.last_frame = None
        super(CameraBufferCleanerThread, self).__init__(name=name)
        self.start()

    def run(self):
        while isRunning:
            ret, self.last_frame = self.camera.read()


siot_config = {
    'client_id': '',
    'server': '192.168.31.233',
    'user': 'siot',
    'password': '*******', # password here
}
siot.init(**siot_config)
siot.connect()
siot.loop()


def ctrl(num):
    assert 0 <= num <= 127
    for i in range(7):
        bit = (num >> i) & 1
        siot.publish('deng/a' + str(6 - i), str(bit))
        siot.publish('deng/b' + str(6 - i), str(bit))
        print(bit)


def get_splitted_img(im):
    _, contour, _ = cv2.findContours(im, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    borders = []
    for i in contour:
        x, y, w, h = cv2.boundingRect(i)
        if np.sum(im[y:y+h+1, x:x+w+1]) >= 100 and np.sum(im[y:y+h+1, x:x+w+1] // 255) / (w * h) < 0.8:
            borders.append(((x, y), (x + w, y + h)))
    borders.sort()

    def resize(im, size=(28, 28)):
        h, w = im.shape
        tar = int(max(h, w) * 1.2)
        up = (tar - h) // 2
        down = tar - up - h
        left = (tar - w) // 2
        right = tar - left - w
        im = np.pad(im, ((up, down), (left, right)), 'constant')
        assert im.shape == (tar, tar)
        return cv2.resize(im, size)
    splitted = []
    for i in borders:
        splitted.append(resize(im[i[0][1]:i[1][1]+1, i[0][0]:i[1][0]+1]))
    return splitted


camera = cv2.VideoCapture(0)
cleaner = CameraBufferCleanerThread(camera)
cv2.namedWindow('Capture')
thr = 127
while isRunning:
    if cleaner.last_frame is None:
        continue
    frame, thr = preproc(cleaner.last_frame, thr) 
    imgs = get_splitted_img(frame)
    key = cv2.waitKey(500) & 0xFF
    if key == ord('q'):
        isRunning = False
        continue
    res = []
    for k, v in enumerate(imgs):
        sock = socket.socket(proto=socket.IPPROTO_TCP)
        sock.connect(('192.168.31.139', 8080))
        msg = ''
        val = v.reshape((1, 784)).squeeze()
        for i in val:
            msg += str(i) + ' '
        sock.send(msg.encode()) # send image data to the host
        prd = int(sock.recv(16).decode().strip())
        sock.close()
        res.append(prd)
    s = ''
    for i in res:
        if i != -1:
            s += chr(i + ord('0'))
    try:
        if s == '':
            print('Null')
            continue
        s = int(s)
        print(s, bin(s))
        ctrl(s)
    except ValueError as e:
        print(e)
        print('s = ' + str(s))

siot.stop()
cleaner.join()
cv2.destroyAllWindows()
